package grpc

import (
	"crypto/tls"
	"crypto/x509"
	"encoding/pem"
	"errors"
	"gitee.com/zxs-micro/zxs-micro-common/grpc/common"
	protocommon "gitee.com/zxs-micro/zxs-micro-common/grpc/protos/common"
	grpc_middleware "github.com/grpc-ecosystem/go-grpc-middleware"
	"google.golang.org/grpc"
	"google.golang.org/grpc/health"
	grpchealth "google.golang.org/grpc/health/grpc_health_v1"
	"net"
	"sync"
	"sync/atomic"
)

type GrpcServer struct {
	lock         *sync.Mutex
	address      string
	listener     net.Listener
	server       *grpc.Server
	healthServer *health.Server
	loadServer   *common.LoadService
	serverCerts  atomic.Value
	tls          *TLSConfig
}

func NewGrpcServer(ctx net.Listener, config GrpcServerConfig) (*GrpcServer, error) {

	return newGrpcServer(ctx, config)
}

func newGrpcServer(ctx net.Listener, config GrpcServerConfig) (*GrpcServer, error) {
	server := new(GrpcServer)
	server.listener = ctx
	server.address = ctx.Addr().String()
	server.lock = &sync.Mutex{}

	var opts = make([]grpc.ServerOption, 0, 0)

	if config.UseTls {
		if config.GrpcSecurity.Cert != nil && config.GrpcSecurity.Key != nil {
			cert, err := tls.X509KeyPair(config.GrpcSecurity.Cert, config.GrpcSecurity.Key)
			if err != nil {
				return nil, err
			}
			server.serverCerts.Store(cert)
			if len(config.GrpcSecurity.Ciphersuite) == 0 {
				config.GrpcSecurity.Ciphersuite = DefaultTLSCipherSuites
			}
			getcert := func(_ *tls.ClientHelloInfo) (*tls.Certificate, error) {
				cert := server.serverCerts.Load().(tls.Certificate)
				return &cert, nil
			}
			server.tls = new(TLSConfig)
			server.tls.config = &tls.Config{
				VerifyPeerCertificate:  config.GrpcSecurity.VerifyCertificate,
				GetCertificate:         getcert,
				SessionTicketsDisabled: true,
				CipherSuites:           config.GrpcSecurity.Ciphersuite,
			}
			server.tls.config.NextProtos = []string{"h2"}
			server.tls.config.MinVersion = tls.VersionTLS12
			server.tls.config.ClientAuth = tls.RequireAndVerifyClientCert
			// if we have client root CAs, create a certPool
			if len(config.ClientRootCA) > 0 {
				server.tls.config.ClientCAs = x509.NewCertPool()
				for _, clientRootCA := range config.ClientRootCA {
					err = server.appendClientRootCA(clientRootCA)
					if err != nil {
						return nil, err
					}
				}
			}
			opts = append(opts, grpc.Creds(&serverCreds{
				serverConfig: server.tls,
			}))
		} else {
			return nil, errors.New("使用tls之前请先配置证书和私钥！")
		}
	}

	// set max send and recv msg sizes
	opts = append(opts, grpc.MaxSendMsgSize(MaxSendMsgSize))
	opts = append(opts, grpc.MaxRecvMsgSize(MaxRecvMsgSize))
	// set the keepalive options
	opts = append(opts, ServerKeepaliveOptions(config.GrpcKeepAlive)...)
	// set connection timeout
	if config.ConnTimeout <= 0 {
		config.ConnTimeout = DefaultConnectionTimeout
	}
	opts = append(opts, grpc.ConnectionTimeout(config.ConnTimeout))
	// set the interceptors
	if len(config.StreamInterceptors) > 0 {
		opts = append(opts, grpc.StreamInterceptor(grpc_middleware.ChainStreamServer(config.StreamInterceptors...)))
	}

	if len(config.UnaryInterceptors) > 0 {
		opts = append(opts, grpc.UnaryInterceptor(grpc_middleware.ChainUnaryServer(config.UnaryInterceptors...)))
	}

	grpcServer := grpc.NewServer(opts...)
	server.server = grpcServer

	server.healthServer = health.NewServer()
	server.healthServer.SetServingStatus("healths", grpchealth.HealthCheckResponse_SERVING)
	grpchealth.RegisterHealthServer(server.server, server.healthServer)

	server.loadServer = common.NewLoadService()
	protocommon.RegisterServerLoadServer(server.server, server.loadServer)

	return server, nil
}

func (s *GrpcServer) appendClientRootCA(clientRoot []byte) error {
	certs, err := pemToX509Certs(clientRoot)
	if err != nil {
		return errors.New("failed to append client root certificate(s): " + err.Error())
	}

	if len(certs) < 1 {
		return errors.New("no client root certificates found")
	}

	for _, cert := range certs {
		s.tls.AddClientRootCA(cert)
	}

	return nil
}

func pemToX509Certs(pemCerts []byte) ([]*x509.Certificate, error) {
	var certs []*x509.Certificate

	// it's possible that multiple certs are encoded
	for len(pemCerts) > 0 {
		var block *pem.Block
		block, pemCerts = pem.Decode(pemCerts)
		if block == nil {
			break
		}

		cert, err := x509.ParseCertificate(block.Bytes)
		if err != nil {
			return nil, err
		}

		certs = append(certs, cert)
	}

	return certs, nil
}

func (s *GrpcServer) Start() error {
	return s.server.Serve(s.listener)
}

func (s *GrpcServer) Stop() {
	s.server.Stop()
}
func (s *GrpcServer) RegisterServer(desc *grpc.ServiceDesc, serv interface{}) {
	s.server.RegisterService(desc, serv)
}

func (s *GrpcServer) Server() *grpc.Server {
	return s.server
}
